// Lasso-LDA
// (c) Yongjin Park, 2013

#include "lda.hh"
#include "lasso.hh"
#include "expression.hh"
#include <string>
#include <iostream>
#include <fstream>

template<>
void output_topics(std::string& output, boost::ptr_vector<func_lasso_gauss_t>& topics)
{
  std::ofstream ofs_beta((output + ".beta").c_str(), std::ios::out);
  std::ofstream ofs_sd((output + ".sd").c_str(), std::ios::out);
  std::ofstream ofs_prec((output + ".prec").c_str(), std::ios::out);
  std::ofstream ofs_mu((output + ".mu").c_str(), std::ios::out);

  for (size_t k = 0; k < topics.size(); ++k)
  {
    const func_lasso_gauss_t& lasso = topics[k];
    lasso.write_beta(ofs_beta);
    ofs_beta << std::endl;
    lasso.write_sd(ofs_sd);
    ofs_sd << std::endl;
    lasso.write_prec(ofs_prec);
    ofs_prec << std::endl;
    lasso.write_mu(ofs_mu);
    ofs_mu << std::endl;
  }
  ofs_beta.close();
  ofs_prec.close();
  ofs_sd.close();
  ofs_mu.close();
}

struct prog_args_t
{

  explicit prog_args_t() :
      g_file(""), x_file(""), y_file(""), K(2), // required args
      output("output"), repeat(1), tol(1e-2), max_iter(100), sample(100), // options
      minibatch(1.)
  {
  }

  std::string g_file;
  std::string x_file;
  std::string y_file;
  size_t K;
  std::string output;
  size_t repeat;
  double tol;
  size_t max_iter;
  size_t sample;
  size_t minibatch;
};

void print_help(const char* prog, const prog_args_t& args)
{
  std::cerr << prog << " -g G.txt -x X.txt -y Y.txt -k K -o out" << std::endl;
  std::cerr << "         -i iter -r repeat -t tol -s sample" << std::endl;
#ifdef ONLINE
  std::cerr << "         -minibatch size" << std::endl;
#endif
  std::cerr << std::endl;
  std::cerr
      << " G.txt  geneset / group file (each line contains row-indexes of X,Y)"
      << std::endl;
  std::cerr << " X.txt  predictor matrix (n x p tab-sep real mat)" << std::endl;
  std::cerr << " Y.txt  response vector (n x 1 real vec)" << std::endl;
  std::cerr << " K      number of topics (default: " << args.K << ")"
      << std::endl;
  std::cerr << " out    output files' header (default: " << args.output << ")"
      << std::endl;
  std::cerr << " iter   max number of iterations (default: " << args.max_iter
      << ")" << std::endl;
  std::cerr << " tol    convergence condition (default: " << args.tol << ")"
      << std::endl;
  std::cerr << " sample number of samples (default: " << args.sample << ")"
      << std::endl;
  std::cerr << " repeat number of repeated fits (default: " << args.repeat
      << ")" << std::endl;
#ifdef ONLINE
  std::cerr << " minibatch size of mini-batch (default: " << args.minibatch
      << ")" << std::endl;
#endif
  std::cerr << std::endl;
}

template<typename T>
T get_param(const char* arg, std::string emsg)
{
  T val;
  try
  {
    val = boost::lexical_cast<T>(arg);
  } catch (boost::bad_lexical_cast&)
  {
    err_msg(emsg);
    exit(1);
  }
  return val;
}

bool parse_args(const int argc, const char* argv[], prog_args_t& args_out)
{
  for (int j = 1; j < argc; ++j)
  {
    std::string curr = argv[j];

    if (curr == "-g" && ++j < argc)
    {
      args_out.g_file = argv[j];
    }
    else if (curr == "-x" && ++j < argc)
    {
      args_out.x_file = argv[j];
    }
    else if (curr == "-y" && ++j < argc)
    {
      args_out.y_file = argv[j];
    }
    else if (curr == "-o" && ++j < argc)
    {
      args_out.output = argv[j];
    }
    else if (curr == "-i" && ++j < argc)
    {
      args_out.max_iter = get_param<size_t>(argv[j], "parsing iter");
    }
    else if (curr == "-minibatch" && ++j < argc)
    {
      args_out.minibatch = get_param<size_t>(argv[j], "parsing minibatch");
    }
    else if (curr == "-s" && ++j < argc)
    {
      args_out.sample = get_param<size_t>(argv[j], "parsing #samples");
    }
    else if (curr == "-k" && ++j < argc)
    {
      args_out.K = get_param<size_t>(argv[j], "parsing K");
    }
    else if (curr == "-r" && ++j < argc)
    {
      args_out.repeat = get_param<size_t>(argv[j], "parsing repeat");
    }
    else if (curr == "-t" && ++j < argc)
    {
      args_out.tol = get_param<double>(argv[j], "parsing tol");
    }
  }

  if (args_out.tol <= 0)
    return false;
  if (args_out.g_file.size() == 0 || args_out.x_file.size() == 0
      || args_out.y_file.size() == 0)
    return false;
  if (args_out.K < 2)
    return false;
  return true;
}

void random_seeding_lasso(std::vector<boost::shared_ptr<expr_lm_t> >& data_vec,
    boost::ptr_vector<expr_ptr_set_t<expr_lm_t> >& data_sets,
    boost::ptr_vector<func_lasso_gauss_t>& topics, const size_t K)
{
  random_seeding<expr_lm_t, func_lasso_gauss_t>(data_vec, data_sets, topics, K);
#ifdef DEBUG
  TLOG("done random seeding")
#endif
}

void write(std::vector<boost::shared_ptr<expr_lm_t> >& data_vec,
    boost::ptr_vector<geneset_t<expr_lm_t> >& genesets,
    boost::ptr_vector<func_lasso_gauss_t>& topics, std::string output)
{
  output_topics<func_lasso_gauss_t>(output, topics);
  output_argmax<expr_lm_t, func_lasso_gauss_t>(output, data_vec, topics);
  output_genesets<expr_lm_t>(output, genesets);

  std::ofstream hyper((output + ".hyper").c_str(), std::ios::out);
  for (size_t j = 0; j < topics.size(); ++j)
  {
    topics[j].write_hyper(hyper);
    hyper << std::endl;
  }
  hyper.close();
}

int main(const int argc, const char* argv[])
{
  prog_args_t args;
  if (!parse_args(argc, argv, args))
  {
    print_help(argv[0], args);
    return -1;
  }

  expr_lm_vec_t data_vec;
  size_t n, p;
  const char* x_file = args.x_file.c_str();
  const char* y_file = args.y_file.c_str();

  if (!read_expr_lm_vec(x_file, y_file, data_vec, n, p))
    return -1;

  // read geneset data
  std::string& g = args.g_file;
  boost::ptr_vector<geneset_t<expr_lm_t> > genesets;
  double a0 = 1. / ((double) args.K);
  const size_t K = args.K;
  if (!read_genesets(g.c_str(), data_vec, a0, K, genesets))
    return -1;

#ifdef ONLINE
  fit_lda_online<expr_lm_t, func_lasso_gauss_t>(data_vec, genesets, p, args);
#else
  fit_lda<expr_lm_t, func_lasso_gauss_t>(data_vec, genesets, p, args);
#endif
  return 0;
}
